import plotly.graph_objects as go
from pymongo import MongoClient
import sys

dbName = sys.argv[1]

# Connect to MongoDB
client = MongoClient('mongodb://localhost:27017')
db = client['cp-test']
# 1k-1m
# collection = db['cp_record_4']
#200-1k
# collection = db['cp_record_5']
collection = db[dbName]

# Retrieve the data
data = collection.find().sort('lengthSource', 1)

# Initialize empty dictionaries for each type
cp_rate_data = {}
cost_time_cp_data = {}
cost_time_dcp_data = {}

# Iterate over the retrieved data
for document in data:
    # Extract the relevant fields from the document
    source_size = document['lengthSe']
    cp_rate = document['cpRate']
    cost_time_cp = document['costTimeCp']
    cost_time_dcp = document['costTimeDcp']
    type = document['type']

    # Append the data to the respective dictionaries based on the type
    cp_rate_data.setdefault(type, []).append((source_size, cp_rate))
    cost_time_cp_data.setdefault(type, []).append((source_size, cost_time_cp))
    cost_time_dcp_data.setdefault(type, []).append((source_size, cost_time_dcp))

# Plot the cpRate graph
fig1 = go.Figure()
for type, data_points in cp_rate_data.items():
    x, y = zip(*data_points)
    fig1.add_trace(go.Scatter(x=x, y=y, mode='lines', name=type))
fig1.update_layout(
    title='Compression Rate',
    xaxis_title='Source Size',
    yaxis_title='Integer cpRate',
)
fig1.show()

# Plot the costTimeCp graph
fig2 = go.Figure()
for type, data_points in cost_time_cp_data.items():
    x, y = zip(*data_points)
    fig2.add_trace(go.Scatter(x=x, y=y, mode='lines', name=type))
fig2.update_layout(
    title='Compression Time',
    xaxis_title='Source Size',
    yaxis_title='Compression Time',
)
fig2.show()

# Plot the costTimeDcp graph
fig3 = go.Figure()
for type, data_points in cost_time_dcp_data.items():
    x, y = zip(*data_points)
    fig3.add_trace(go.Scatter(x=x, y=y, mode='lines', name=type))
fig3.update_layout(
    title='Decompression Time',
    xaxis_title='Source Size',
    yaxis_title='Decompression Time',
)
fig3.show()